# coding = utf-8

import gradio as gr

from langchain_openai import OpenAIEmbeddings # type: ignore
from langchain.chains import RetrievalQA # type: ignore
from langchain_openai import ChatOpenAI # type: ignore
from langchain.document_loaders import TextLoader # type: ignore
from langchain_community.vectorstores import FAISS # type: ignore
from langchain.text_splitter import CharacterTextSplitter # type: ignore
import os
os.environ['PYTHONUTF8'] = "1"
os.environ['PYTHONIOENCODING'] = 'gbk'

def initialize_sales_bot(document_path: str):
    ##db = FAISS.load_local(vector_store_dir, OpenAIEmbeddings(),allow_dangerous_deserialization=True)
    db = init_db(document_path=document_path)
    llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0, verbose=True)

 
    global SALES_BOT
    SALES_BOT = RetrievalQA.from_chain_type(
        llm,
        retriever=db.as_retriever(
            search_type="similarity_score_threshold",
            search_kwargs={"score_threshold": 0.8},
        ),
    )
    # 返回向量数据库的检索结果
    SALES_BOT.return_source_documents = True
    
    print(SALES_BOT)

    return SALES_BOT


def sales_chat(message, history):
    initial_messages = [
        ("系统", "请问有什么我可以帮您？"),
    ]
    print(f"[message]{message}")
    print(f"[history]{history}")
    # TODO: 从命令行参数中获取
    enable_chat = True

    ans = SALES_BOT({"query": message})
    # 如果检索出结果，或者开了大模型聊天模式
    # 返回 RetrievalQA combine_documents_chain 整合的结果
    if ans["source_documents"] or enable_chat:
        print(f"[result]{ans['result']}")
        print(f"[source_documents]{ans['source_documents']}")
        return ans["result"]
    # 否则输出套路话术
    else:
        return "这个问题我要问问领导"


def launch_gradio():
    demo = gr.ChatInterface(
        fn=sales_chat,
        title="智能客服",
        examples=["这手机保修期多久?", "这手机防水吗?", "我能不能试用产品?"],
        retry_btn=None,
        undo_btn=None,
        chatbot=gr.Chatbot(height=600),
    )

    demo.launch(share=True, server_name="0.0.0.0")

def init_db(vector_store_dir: str = "real_estates_sale",document_path: str="real_estate_sales_data.txt"):
    """"""
    loader = TextLoader(document_path)
    # 加载文档
    documents = loader.load()
    
    # 实例化文本分割器
    text_splitter = CharacterTextSplitter(        
        separator = r'\d+\.',
        chunk_size = 100,
        chunk_overlap  = 0,
        length_function = len,
        is_separator_regex = True,
    )
    # 分割文本
    docs = text_splitter.split_documents(documents)
    
    for doc in docs:
        print(doc)
    # FAISS 向量数据库，使用 docs 的向量作为初始化存储
    db = FAISS.from_documents(docs, OpenAIEmbeddings())
    return db

if __name__ == "__main__":

    # 初始化房产销售机器人
    #initialize_sales_bot(document_path="real_estate_sales_data.txt")
    initialize_sales_bot(document_path="手机销售语料库.txt")

    # 启动 Gradio 服务
    launch_gradio()

